#!/usr/bin/env python3
"""
模型文件加密工具 - V4.0零信任架构
用于加密PyTorch模型文件，保护AI知识产权

使用方法:
python encrypt_model.py

功能:
1. 加密大型模型文件(支持分块处理)
2. 生成完整性校验码
3. 安全内存管理
4. 加密过程验证

⚠️ 安全警告:
- 加密后请删除原始模型文件
- 密钥必须安全存储
- 定期轮换加密密钥
"""

import os
import sys
import hashlib
import base64
import time
from pathlib import Path
from typing import Tuple, Optional
from cryptography.fernet import Fernet


class ModelEncryptor:
    """模型文件加密器 - 支持大文件分块加密"""

    def __init__(self, key: bytes):
        """初始化加密器"""
        if not self._validate_key(key):
            raise ValueError("无效的加密密钥")
        self.fernet = Fernet(key)
        self.chunk_size = 1024 * 1024  # 1MB分块大小

    def _validate_key(self, key: bytes) -> bool:
        """验证密钥格式"""
        try:
            Fernet(key)
            return True
        except Exception:
            return False

    def encrypt_file(self, input_path: str, output_path: str) -> Tuple[bool, str]:
        """
        加密文件

        Args:
            input_path: 输入文件路径
            output_path: 输出文件路径

        Returns:
            (success, message)
        """
        try:
            input_file = Path(input_path)
            output_file = Path(output_path)

            # 验证输入文件
            if not input_file.exists():
                return False, f"输入文件不存在: {input_path}"

            if not input_file.is_file():
                return False, f"输入路径不是文件: {input_path}"

            # 计算文件大小
            file_size = input_file.stat().st_size
            print(f"📁 加密文件: {input_path}")
            print(f"📊 文件大小: {file_size / (1024*1024):.1f} MB")

            # 生成文件完整性校验码
            print("🔐 计算原始文件校验码...")
            original_hash = self._calculate_file_hash(input_path)

            # 开始加密
            print(f"🔒 开始加密，分块大小: {self.chunk_size / (1024*1024):.1f} MB")
            start_time = time.time()

            success, message = self._encrypt_large_file(input_path, output_path)

            if not success:
                return False, message

            # 验证加密结果
            print("✅ 验证加密结果...")
            if not self._verify_encryption(output_path):
                return False, "加密结果验证失败"

            encryption_time = time.time() - start_time
            print(f"⏱️  加密完成，耗时: {encryption_time:.2f} 秒")

            # 生成元数据文件
            metadata = {
                'original_filename': input_file.name,
                'original_size': file_size,
                'original_hash': original_hash,
                'encrypted_size': output_file.stat().st_size,
                'encryption_time': encryption_time,
                'chunk_size': self.chunk_size,
                'encryption_algorithm': 'AES-256-CBC (Fernet)',
                'timestamp': time.time()
            }

            metadata_path = output_path.with_suffix('.metadata')
            self._save_metadata(metadata, metadata_path)

            print(f"📋 元数据已保存: {metadata_path}")
            print(f"✅ 加密完成: {output_path}")

            return True, "加密成功"

        except Exception as e:
            return False, f"加密过程出错: {str(e)}"

    def _encrypt_large_file(self, input_path: Path, output_path: Path) -> Tuple[bool, str]:
        """加密大文件（分块处理）"""
        try:
            with open(input_path, 'rb') as infile:
                with open(output_path, 'wb') as outfile:
                    chunk_count = 0

                    while True:
                        chunk = infile.read(self.chunk_size)
                        if not chunk:
                            break

                        # 加密分块
                        encrypted_chunk = self.fernet.encrypt(chunk)
                        outfile.write(encrypted_chunk)

                        chunk_count += 1
                        # 显示进度
                        if chunk_count % 10 == 0:
                            processed_mb = (chunk_count * self.chunk_size) / (1024 * 1024)
                            print(f"   处理进度: {processed_mb:.1f} MB")

            return True, "加密完成"

        except Exception as e:
            return False, f"文件加密失败: {str(e)}"

    def _calculate_file_hash(self, file_path: Path) -> str:
        """计算文件SHA-256哈希"""
        sha256_hash = hashlib.sha256()
        with open(file_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                sha256_hash.update(chunk)
        return sha256_hash.hexdigest()

    def _verify_encryption(self, encrypted_path: Path) -> bool:
        """验证加密文件可以被解密"""
        try:
            with open(encrypted_path, 'rb') as f:
                # 只验证第一个分块
                first_chunk = f.read(self.chunk_size)
                if first_chunk:
                    self.fernet.decrypt(first_chunk)
            return True
        except Exception:
            return False

    def _save_metadata(self, metadata: dict, metadata_path: Path):
        """保存元数据"""
        import json
        with open(metadata_path, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, indent=2, ensure_ascii=False)


def get_encryption_key() -> bytes:
    """获取加密密钥"""
    print("🔑 获取加密密钥")

    # 尝试从环境变量获取
    key_env = os.getenv('MODEL_ENCRYPTION_KEY')
    if key_env:
        try:
            key = key_env.encode('utf-8')
            # 验证密钥格式
            Fernet(key)
            print("✅ 从环境变量获取密钥成功")
            return key
        except Exception:
            print("❌ 环境变量中的密钥格式无效")

    # 尝试从文件获取
    key_file = Path("model_encryption.key")
    if key_file.exists():
        try:
            with open(key_file, 'rb') as f:
                key = f.read()
            # 验证密钥格式
            Fernet(key)
            print(f"✅ 从文件获取密钥成功: {key_file}")
            return key
        except Exception:
            print("❌ 密钥文件格式无效")

    # 手动输入密钥
    while True:
        try:
            print("\n请输入AES-256密钥 (Base64编码，44字符):")
            key_input = input().strip()
            key = key_input.encode('utf-8')

            # 验证密钥格式
            Fernet(key)
            print("✅ 密钥格式验证通过")
            return key

        except KeyboardInterrupt:
            print("\n\n❌ 操作已取消")
            sys.exit(1)
        except Exception as e:
            print(f"❌ 密钥格式错误: {e}")
            print("请确保密钥是44字符的Base64编码字符串")


def find_model_file() -> Optional[str]:
    """查找模型文件"""
    possible_paths = [
        "uc_model/final_weights/uceis_best_model.pth",
        "../uc_model/final_weights/uceis_best_model.pth",
        "uceis_best_model.pth"
    ]

    for path in possible_paths:
        if Path(path).exists():
            return path

    print("❌ 未找到模型文件")
    print("请确保以下文件之一存在:")
    for path in possible_paths:
        print(f"   - {path}")

    return None


def display_encryption_info():
    """显示加密信息"""
    print("=" * 70)
    print("🔐 模型文件加密工具 - V4.0零信任架构")
    print("=" * 70)
    print("📋 功能说明:")
    print("   • AES-256-CBC加密算法")
    print("   • 分块处理支持大文件")
    print("   • 文件完整性校验")
    print("   • 安全内存管理")
    print("=" * 70)
    print("⚠️  安全提醒:")
    print("   1. 加密完成后请删除原始模型文件")
    print("   2. 密钥必须安全存储")
    print("   3. 加密文件请妥善保管")
    print("=" * 70)


def main():
    """主函数"""
    display_encryption_info()

    try:
        # 获取加密密钥
        key = get_encryption_key()

        # 查找模型文件
        model_path = find_model_file()
        if not model_path:
            sys.exit(1)

        # 确认加密操作
        print(f"\n📄 准备加密模型文件: {model_path}")
        try:
            confirm = input("确认继续加密? (y/n): ").strip().lower()
            if confirm not in ['y', 'yes', '是']:
                print("❌ 操作已取消")
                sys.exit(0)
        except KeyboardInterrupt:
            print("\n❌ 操作已取消")
            sys.exit(0)

        # 创建加密器
        encryptor = ModelEncryptor(key)

        # 确定输出路径
        model_file = Path(model_path)
        encrypted_path = model_file.with_suffix('.pth.enc')

        print(f"📁 输出路径: {encrypted_path}")

        # 执行加密
        success, message = encryptor.encrypt_file(model_path, encrypted_path)

        if success:
            print("\n🎉 模型加密成功!")
            print(f"📁 加密文件: {encrypted_path}")
            print(f"📊 元数据文件: {encrypted_path.with_suffix('.metadata')}")

            # 询问是否删除原文件
            try:
                delete_original = input("\n是否删除原始模型文件? (y/n): ").strip().lower()
                if delete_original in ['y', 'yes', '是']:
                    os.remove(model_path)
                    print(f"✅ 原始文件已删除: {model_path}")
                    print("⚠️  请确保已安全保存加密文件和密钥")
            except KeyboardInterrupt:
                print("\n跳过删除原始文件")

            print("\n📝 下一步:")
            print("   1. 更新配置文件中的模型路径")
            print("   2. 在生产环境设置 MODEL_ENCRYPTION_KEY")
            print("   3. 启动V4.0安全推理服务")
        else:
            print(f"\n❌ 加密失败: {message}")
            sys.exit(1)

    except KeyboardInterrupt:
        print("\n\n❌ 操作已取消")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 加密过程出错: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()